import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
from matplotlib import pyplot as plt
import statsmodels.api as sm
import seaborn as sns
import numpy as np
# 标准化数据（主成分分析前通常需要标准化）
from sklearn.preprocessing import StandardScaler
scaler = StandardScaler()

db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': 'sjk123',
    'database': 'tushare',
    'port': 3306,
    'charset': 'utf8mb4'
}
engine = create_engine(f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}")
conn = engine.connect()
chunk_size = 10000

df = pd.read_sql_query("""
                     SELECT d.*,m.buy_sm_vol, m.sell_sm_vol,m.buy_md_vol, m.sell_md_vol,m.buy_lg_vol, m.sell_lg_vol,m.buy_elg_vol, m.sell_elg_vol, m.net_mf_vol,i.vol as i_vol, i.closes as i_closes  
                     FROM date_1 d 
                     JOIN moneyflows m ON d.ts_code = m.ts_code AND d.trade_date = m.trade_date 
                     LEFT JOIN index_daily i on d.trade_date = i.trade_date AND i.ts_code = '399001.SZ' 
                     WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' AND d.ts_code = '002229.SZ'""",
                     conn,
                     chunksize=chunk_size)

df1 = pd.concat(df, ignore_index=True)
print(df1.head())

# 修正涨跌幅计算公式 - 添加括号
df1['zd_closes'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 4)
df1['zs_closes'] = round((df1['i_closes'] - df1['i_closes'].shift(1)) / df1['i_closes'].shift(1), 4)
df1['zs_vol'] = round((df1['i_vol'] - df1['i_vol'].shift(1)) / df1['i_vol'].shift(1), 4)

# 修正股票的涨跌幅计算 - 使用正确的公式
df1['zd_pct_chg'] = round((df1['pct_chg'] - df1['pct_chg'].shift(1)) / df1['pct_chg'].shift(1), 4)

print(df1.head())

# 处理缺失值
df1 = df1.dropna(subset=['zd_pct_chg']).reset_index(drop=True)

# 选择需要进行主成分分析的自变量
X = df1[['vol', 'amount', 'buy_sm_vol', 'sell_sm_vol', 'buy_md_vol', 'sell_md_vol', 
         'buy_lg_vol', 'sell_lg_vol', "buy_elg_vol", "sell_elg_vol"]]


X_scaled = scaler.fit_transform(X)

# 计算特征值和特征向量 - 使用标准化后的数据
cov_matrix = np.cov(X_scaled, rowvar=False)
eigenvalues, eigenvectors = np.linalg.eig(cov_matrix)

# 按特征值大小排序
idx = eigenvalues.argsort()[::-1]   
eigenvalues = eigenvalues[idx]
eigenvectors = eigenvectors[:, idx]

print('累计共享率：', round(eigenvalues[:5].sum() / eigenvalues.sum(), 2) * 100, '%')

# 选择要保留的主成分个数
n_components = 5
top_eigenvectors = eigenvectors[:, :n_components]

# 计算主成分
principal_components = np.dot(X_scaled, top_eigenvectors)

# 将主成分添加到原数据中
data_pca = pd.DataFrame(principal_components, columns=[f'PC{i+1}' for i in range(n_components)])

# 确保y的索引与X_pca一致 - 修正：使用df1中的zd_pct_chg
y = df1['zd_pct_chg'].copy().iloc[:len(data_pca)]  # 确保长度一致

# 构建回归模型
X_pca = data_pca[[f'PC{i+1}' for i in range(n_components)]].copy()
X_pca = sm.add_constant(X_pca)  # 添加常数项

model = sm.OLS(y, X_pca)
results = model.fit()

# 输出结果
print('回归模型结果：')
print(results.summary())

# 选取pc3、pc4、pc5作为新的自变量
X_pca_selected = data_pca[['PC3', 'PC4', 'PC5']].copy()
X_pca_selected.columns = ['PC3', 'PC4', 'PC5']

# 添加常数项
X_pca_selected = sm.add_constant(X_pca_selected)

# 构建回归模型
model_selected = sm.OLS(y, X_pca_selected)
results_selected = model_selected.fit()

# 输出结果
print('选取PC3、PC4、PC5的回归模型结果：')
print(results_selected.summary())

# 绘制散点图
fig, axes = plt.subplots(1, 3, figsize=(15, 5))  # 修正：3个子图对应3个主成分
for i, col in enumerate(['PC3', 'PC4', 'PC5']):
    axes[i].scatter(data_pca[col], y, s=50, alpha=0.5)
    axes[i].set_xlabel(col)
    axes[i].set_ylabel('zd_pct_chg')
    axes[i].set_title(f'{col} vs zd_pct_chg')
plt.tight_layout()
plt.show()

# 计算相关系数并输出主成分表达式
print("\n主成分表达式：")
for k in range(5):
    string_y = f'PC{k+1} = '
    i = eigenvectors[:, k]  # 修正：获取第k个特征向量
    
    for j in range(len(i)):
        coef = round(i[j], 2)
        if j == 0:
            if coef >= 0:
                string_y += f'{coef} * X_{j+1}'
            else:
                string_y += f'-{abs(coef)} * X_{j+1}'
        else:
            if coef >= 0:
                string_y += f' + {coef} * X_{j+1}'
            else:
                string_y += f' - {abs(coef)} * X_{j+1}'
    
    if k != 2 and k != 4:  # 按照原代码逻辑，跳过第3和第5个主成分
        print(string_y)

# 关闭数据库连接
conn.close()